package geppetto.cat.alignments.phrasal;

import java.io.PrintStream;

import geppetto.cat.alignments.Alignment;
import geppetto.cat.constrains.ConstrainedProjectionStats;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.stats.EStepStats;
import geppetto.phraseHMM.WordTrie;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import geppetto.phraseTable.phrase.Phrase;
import gnu.trove.TIntDoubleHashMap;
import gnu.trove.TIntObjectHashMap;

public class PhrasalHmmAgreementModel extends AbstractModel{

	PhrasalHmmModel directModel;
	PhrasalHmmModel reverseModel;
	WordTrie sourceTrie;
	WordTrie targetTrie;
	int numberOfPhrasesInCache = 1000;

	public PhrasalHmmAgreementModel(PhrasalHmmModel forwardModel,
			PhrasalHmmModel backwardModel) {
		super();
		this._corpus = forwardModel._corpus;
		this.directModel = forwardModel;
		sourceTrie = directModel.sourceTrie;
		targetTrie = directModel.targetTrie;
		this.reverseModel = backwardModel;
	}

	@Override
	public EStepStats eStep() {
		clearTables();
		EStepStats stats = createModelStats();
		byte sentenceSource = BilingualCorpus.TRAIN_CORPUS;
		int numberOfSentences = _corpus.getNumSentences(sentenceSource);
		int numberOfSentencesPerPrint = (int)Math.pow(10.0d, String.valueOf(numberOfSentences).length()-1);
		for (int i = 0; i < numberOfSentences; i++) {
			if(i % numberOfSentencesPerPrint == 0){
				System.err.println("processing sentences " + i + " to " + (i+numberOfSentencesPerPrint));
			}
			directModel.generatePosteriors(i, sentenceSource);
			reverseModel.generatePosteriors(i, sentenceSource);			
			agreement(i, sentenceSource);
			directModel.updateExpectedCounts();
			reverseModel.updateExpectedCounts();
		}
		return stats;
	}
	
	public void extractPhrases(PrintStream phraseExtractionFile, PrintStream nullPhraseExtractionFile, PrintStream reorderingExtractionFile, PrintStream phraseExtractionFileB, PrintStream nullPhraseExtractionFileB, PrintStream reorderingExtractionFileB){
				clearTables();
		byte sentenceSource = BilingualCorpus.TRAIN_CORPUS;
		int numberOfSentences = _corpus.getNumSentences(sentenceSource);
		int numberOfSentencesPerPrint = (int)Math.pow(10.0d, String.valueOf(numberOfSentences).length()-1);
		for (int i = 0; i < numberOfSentences; i++) {
			//if(i % numberOfSentencesPerPrint == 0){
				System.err.println("processing sentences " + i + " to " + (i+numberOfSentencesPerPrint));
			//}
			directModel.generatePosteriors(i, sentenceSource);
			reverseModel.generatePosteriors(i, sentenceSource);
			agreement(i, sentenceSource);
			if(i % numberOfPhrasesInCache == 0){
				directModel.printPhrasePairExpectations(phraseExtractionFile, nullPhraseExtractionFile);
				reverseModel.printPhrasePairExpectations(phraseExtractionFileB, nullPhraseExtractionFileB);
				directModel.printReorderingExpectations(reorderingExtractionFile);
				reverseModel.printReorderingExpectations(reorderingExtractionFileB);
				directModel.clearTables();
				reverseModel.clearTables();
			}
			directModel.updateExpectedCounts();
			reverseModel.updateExpectedCounts();
		}
		directModel.printPhrasePairExpectations(phraseExtractionFile, nullPhraseExtractionFile);
		reverseModel.printPhrasePairExpectations(phraseExtractionFileB, nullPhraseExtractionFileB);
		directModel.printReorderingExpectations(reorderingExtractionFile);
		reverseModel.printReorderingExpectations(reorderingExtractionFileB);
	}
	
	public void writeModels(PrintStream phraseTableFile, PrintStream nullPhraseTableFile, PrintStream reorderingTableFile, PrintStream phraseTableFileB, PrintStream nullPhraseTableFileB, PrintStream reorderingTableFileB){
		mStep();
		directModel.printPhraseTable(phraseTableFile, nullPhraseTableFile);
		reverseModel.printPhraseTable(phraseTableFileB, nullPhraseTableFileB);
		directModel.printReorderingTable(reorderingTableFile);
		reverseModel.printReorderingTable(reorderingTableFileB);
	}
	
	public void clearTables(){
		directModel.clearTables();
		reverseModel.clearTables();
	}

	@Override
	public void finalizeTrain() {
		// TODO Auto-generated method stub

	}

	@Override
	public String getName() {		
		return "agreement";
	}

	@Override
	public double getNullPhrasePosterior(int sentenceNumber,
			byte sentenceSource, int[] foreingSentence, int[] sourceSentence,
			int startForeignIndex, int endForeignIndex) {
		// TODO Auto-generated method stub
		return 0;
	}

	@Override
	public double getPhrasePosterior(int sentenceNumber, byte sentenceSource,
			int[] foreingSentence, int[] sourceSentence, int startSourceIndex,
			int endSourceIndex, int startForeignIndex, int endForeignIndex) {
		// TODO Auto-generated method stub
		return 0;
	}

	@Override
	public void initializeTrain() {
		directModel.initializeTrain();
		reverseModel.initializeTrain();
	}

	@Override
	public void mStep() {
		directModel.mStep();
		reverseModel.mStep();
	}

	@Override
	public Alignment posteriorDecodingAlignment(int sentenceNumber,
			byte sentenceSource, float treshhold, boolean projectPosteriors,
			ConstrainedProjectionStats stats) {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public Alignment viterbiAlignment(int sentenceNumber, byte sentenceSource,
			boolean projectPosteriors, ConstrainedProjectionStats stats) {
		// TODO Auto-generated method stub
		return null;
	}


	public void agreement(int sentenceNumber, byte sentenceSource){
		int[] source = _corpus.getSourceSentence(sentenceNumber, sentenceSource);
		int[] target = _corpus.getForeignSentence(sentenceNumber, sentenceSource);
		ExtractedPhrasePair[] directModelPairs = directModel.extractedPhrasesFromSentenceInEstep;
		ExtractedPhrasePair[] reverseModelPairs = reverseModel.extractedPhrasesFromSentenceInEstep;
		ExtractedPhrasePair[][][][] directPhraseReference = new ExtractedPhrasePair[source.length][source.length][target.length][target.length];
		for(ExtractedPhrasePair pp : directModelPairs){
			if(!pp.isNullPhrase()){
				directPhraseReference[pp.getSourceStartPos()][pp.getSourceEndPos()][pp.getForeignStartPos()][pp.getForeignEndPos()]=pp;
			}
		}
		for(ExtractedPhrasePair rpp : reverseModelPairs){
			if(!rpp.isNullPhrase()){
				ExtractedPhrasePair pp = directPhraseReference[rpp.getForeignStartPos()][rpp.getForeignEndPos()][rpp.getSourceStartPos()][rpp.getSourceEndPos()];
				double agreementScore = pp.getScore() * rpp.getScore();
				pp.setScore(agreementScore);
				rpp.setScore(agreementScore);
			}
		}
	}

	public void printPhraseTablesForward(PrintStream out, double threshold){
		PhraseTranslationTable directPhraseTable = directModel.translationTable;
		PhraseTranslationTable reversePhraseTable = reverseModel.translationTable;
		TIntObjectHashMap<TIntDoubleHashMap> directPhrases = directPhraseTable.getPhraseScores();
		TIntObjectHashMap<TIntDoubleHashMap> reversePhrases = reversePhraseTable.getPhraseScores();
		for(int s : directPhrases.keys()){
			TIntDoubleHashMap directTargetPhrases = directPhrases.get(s); 
			for(int t : directTargetPhrases.keys()){
				double score = directTargetPhrases.get(t);
				double reverseScore = reversePhrases.get(t).get(s);
				if(score > threshold || reverseScore > threshold){
					out.println(Phrase.getSourceString(_corpus, sourceTrie.getPhraseIndexesById(s)) + " ||| " + Phrase.getTargetString(_corpus, targetTrie.getPhraseIndexesById(t)) + " ||| " + score + " " + reverseScore);
				}
			}
		}
	}
}
